package rs.pnv.util.xmlsec;

import java.io.BufferedInputStream;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLConnection;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Collection;
import java.util.Iterator;

import javax.xml.XMLConstants;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.FactoryConfigurationError;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.soap.SOAPBody;
import javax.xml.soap.SOAPPart;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.Source;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerConfigurationException;
import javax.xml.transform.TransformerException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.TransformerFactoryConfigurationError;
import javax.xml.transform.dom.DOMResult;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;

import org.apache.xml.security.exceptions.XMLSecurityException;
import org.apache.xml.security.keys.KeyInfo;
import org.apache.xml.security.keys.keyresolver.implementations.RSAKeyValueResolver;
import org.apache.xml.security.keys.keyresolver.implementations.X509CertificateResolver;
import org.apache.xml.security.signature.XMLSignature;
import org.apache.xml.security.signature.XMLSignatureException;
import org.apache.xml.security.transforms.TransformationException;
import org.apache.xml.security.transforms.Transforms;
import org.apache.xml.security.transforms.params.XPathContainer;
import org.apache.xml.security.utils.Constants;
import org.w3c.dom.DOMException;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;
import org.xml.sax.SAXException;

public class SignDocUtil {

	public static Document removeSignatureFromXML(Document doc) {

		Element element = (Element) doc.getElementsByTagNameNS(
				"http://www.w3.org/2000/09/xmldsig#", "Signature").item(0);
		if (element != null) {
			element.getParentNode().removeChild(element);
		}
		return doc;
	}

	public static X509Certificate getCertificate(Document doc) {
		try {
			// Pronalazi se prvi Signature element
			NodeList signatures = doc.getElementsByTagNameNS(
					"http://www.w3.org/2000/09/xmldsig#", "Signature");
			Element signatureEl = (Element) signatures.item(0);

			// kreira se signature objekat od elementa
			XMLSignature signature = new XMLSignature(signatureEl, null);
			// preuzima se key info
			KeyInfo keyInfo = signature.getKeyInfo();
			// ako postoji
			if (keyInfo != null) {
				// registruju se resolver-i za javni kljuc i sertifikat
				keyInfo.registerInternalKeyResolver(new RSAKeyValueResolver());
				keyInfo.registerInternalKeyResolver(new X509CertificateResolver());

				// ako sadrzi sertifikat
				if (keyInfo.containsX509Data()
						&& keyInfo.itemX509Data(0).containsCertificate()) {
					X509Certificate cert = keyInfo.itemX509Data(0)
							.itemCertificate(0).getX509Certificate();
					// ako postoji sertifikat, provera potpisa
					if (cert != null)
						return cert;
					else
						return null;
				} else
					return null;
			} else
				return null;

		} catch (XMLSignatureException e) {
			e.printStackTrace();
			return null;
		} catch (XMLSecurityException e) {
			e.printStackTrace();
			return null;
		}
	}

	public static boolean verifySignature(Document doc) {

		try {
			// Pronalazi se prvi Signature element
			NodeList signatures = doc.getElementsByTagNameNS(
					"http://www.w3.org/2000/09/xmldsig#", "Signature");
			Element signatureEl = (Element) signatures.item(0);

			// kreira se signature objekat od elementa
			XMLSignature signature = new XMLSignature(signatureEl, null);
			// preuzima se key info
			KeyInfo keyInfo = signature.getKeyInfo();
			// ako postoji
			if (keyInfo != null) {
				// registruju se resolver-i za javni kljuc i sertifikat
				keyInfo.registerInternalKeyResolver(new RSAKeyValueResolver());
				keyInfo.registerInternalKeyResolver(new X509CertificateResolver());

				// ako sadrzi sertifikat
				if (keyInfo.containsX509Data()
						&& keyInfo.itemX509Data(0).containsCertificate()) {
					Certificate cert = keyInfo.itemX509Data(0)
							.itemCertificate(0).getX509Certificate();
					// ako postoji sertifikat, provera potpisa
					if (cert != null)
						return signature
								.checkSignatureValue((X509Certificate) cert);
					else
						return false;
				} else
					return false;
			} else
				return false;

		} catch (XMLSignatureException e) {
			e.printStackTrace();
			return false;
		} catch (XMLSecurityException e) {
			e.printStackTrace();
			return false;
		}
	}
	
	

	/**
	 * Kreira DOM od XML dokumenta
	 */
	public static Document loadDocument(String file) {
		try {
			DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
			dbf.setNamespaceAware(true);
			DocumentBuilder db = dbf.newDocumentBuilder();
			Document document = db.parse(new File(file));

			return document;
		} catch (FactoryConfigurationError e) {
			e.printStackTrace();
			return null;
		} catch (ParserConfigurationException e) {
			e.printStackTrace();
			return null;
		} catch (SAXException e) {
			e.printStackTrace();
			return null;
		} catch (IOException e) {
			e.printStackTrace();
			return null;
		}
	}

	@SuppressWarnings("rawtypes")
	public static Certificate readRemoteCertificate(String cerPath) {
		try {

			URL url = new URL(cerPath);
			URLConnection connection = url.openConnection();
			CertificateFactory cf = CertificateFactory.getInstance("X.509");
			Collection c = cf.generateCertificates(connection.getInputStream());
			Iterator i = c.iterator();
			while (i.hasNext()) {
				Certificate cert = (Certificate) i.next();
				return cert;
			}
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (CertificateException e) {
			e.printStackTrace();
		} catch (MalformedURLException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
		return null;
	}

	public static Certificate readCertificate(String pathToKeyStore,
			String password) {
		try {
			// kreiramo instancu KeyStore
			KeyStore ks = KeyStore.getInstance("JKS", "SUN");
			// ucitavamo podatke
			BufferedInputStream in = new BufferedInputStream(
					new FileInputStream(pathToKeyStore));
			ks.load(in, password.toCharArray());

			if (ks.isKeyEntry(password)) {
				Certificate cert = ks.getCertificate(password);
				return cert;

			} else
				return null;

		} catch (KeyStoreException e) {
			e.printStackTrace();
			return null;
		} catch (NoSuchProviderException e) {
			e.printStackTrace();
			return null;
		} catch (FileNotFoundException e) {
			e.printStackTrace();
			return null;
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
			return null;
		} catch (CertificateException e) {
			e.printStackTrace();
			return null;
		} catch (IOException e) {
			e.printStackTrace();
			return null;
		}
	}

	public static Document reserialize(Node request) {
		Document r = null;
		try {
			// potrebno je napisati metodu getDocumentBuilder, koja kreira
			// DocumentBuilder
			// documentBuilder mora imati namespaceAware = true
			DocumentBuilder db = getDocumentBuilder();
			Transformer transformer = TransformerFactory.newInstance()
					.newTransformer();
			transformer.setOutputProperty(OutputKeys.INDENT, "no");
			ByteArrayOutputStream baos = new ByteArrayOutputStream();
			StreamResult result = new StreamResult(baos);
			DOMSource src = new DOMSource(request);
			transformer.transform(src, result);
			ByteArrayInputStream bais = new ByteArrayInputStream(
					baos.toByteArray());
			r = db.parse(bais);
		} catch (Exception e) {
			e.printStackTrace();
		}
		return r;
	}

	public static DocumentBuilder getDocumentBuilder() {
		try {
			DocumentBuilderFactory docBuilderFactory = DocumentBuilderFactory
					.newInstance();
			docBuilderFactory.setNamespaceAware(true);
			DocumentBuilder builder = docBuilderFactory.newDocumentBuilder();
			return builder;
		} catch (ParserConfigurationException e) {
			e.printStackTrace();
			return null;
		}
	}

	public static Document signDocument(Document doc, PrivateKey privateKey,Certificate cert) {
			
		
		
		try {
			Element rootEl = doc.getDocumentElement();
			rootEl.setAttributeNS(XMLConstants.XMLNS_ATTRIBUTE_NS_URI, "xmlns:soap", "http://schemas.xmlsoap.org/soap/envelope/");
			// kreira se signature objekat
			XMLSignature sig = new XMLSignature(doc, null,
					XMLSignature.ALGO_ID_SIGNATURE_RSA_SHA1);
			// kreiraju se transformacije nad dokumentom
			Transforms transforms = new Transforms(doc);
			
			// iz potpisa uklanja Signature element
			// Ovo je potrebno za enveloped tip po specifikaciji
			transforms.addTransform(Transforms.TRANSFORM_ENVELOPED_SIGNATURE);
			// normalizacija
			transforms.addTransform(Transforms.TRANSFORM_C14N_WITH_COMMENTS);

			// potpisuje se citav dokument (URI "")
			sig.addDocument("", transforms, Constants.ALGO_ID_DIGEST_SHA1);

			// U KeyInfo se postavalja Javni kljuc samostalno i citav sertifikat
			sig.addKeyInfo(cert.getPublicKey());
			sig.addKeyInfo((X509Certificate) cert);

			// poptis je child root elementa
			rootEl.appendChild(sig.getElement());

			// potpisivanje
			sig.sign(privateKey);

			return doc;

		} catch (TransformationException e) {
			e.printStackTrace();
			return null;
		} catch (XMLSignatureException e) {
			e.printStackTrace();
			return null;
		} catch (DOMException e) {
			e.printStackTrace();
			return null;
		} catch (XMLSecurityException e) {
			e.printStackTrace();
			return null;
		}
	}

	/**
	 * Snima DOM u XML fajl
	 */
	public static void saveDocument(Document doc, String fileName) {
		try {
			File outFile = new File(fileName);
			FileOutputStream f = new FileOutputStream(outFile);

			TransformerFactory factory = TransformerFactory.newInstance();
			Transformer transformer = factory.newTransformer();

			DOMSource source = new DOMSource(doc);
			StreamResult result = new StreamResult(f);

			transformer.transform(source, result);

			f.close();

		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (TransformerConfigurationException e) {
			e.printStackTrace();
		} catch (IllegalArgumentException e) {
			e.printStackTrace();
		} catch (MalformedURLException e) {
			e.printStackTrace();
		} catch (TransformerFactoryConfigurationError e) {
			e.printStackTrace();
		} catch (TransformerException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	/**
	 * 
	 * @param doc
	 *            - koji se kriptuje
	 * @param jksLocation
	 *            - lokacija jksa
	 * @param jksPassword
	 *            - lozinka
	 * @return
	 */
	public Document docSignature(Document doc, String jksLocation,
			String jksPassword, String cerPath) {
		Decrypt decrypt = new Decrypt();
		System.out.println("*****Potpisivanje*****");
		System.out.println("Ucitava se privatni kljuc BankeA...");
		PrivateKey privateKey = decrypt
				.readPrivateKey(jksLocation, jksPassword);
		System.out.println("Ucitavanje setifikata...");
		// Certificate certificate = readRemoteCertificate(cerPath) ;
		Certificate certificate = readCerFromFile(cerPath);
		System.out.println("Reserializacija...");
		doc = reserialize(doc);
		System.out.println("Potpisivanje...");
		Document retDoc = signDocument(doc, privateKey, certificate);
		System.out.println("*****Potpisivanje*****");
		return retDoc;
	}
	
	public void bodySignature(SOAPPart body, String jksLocation,
			String jksPassword, String cerPath) {
		Decrypt decrypt = new Decrypt();
		System.out.println("*****Potpisivanje*****");
		System.out.println("Ucitava se privatni kljuc BankeA...");
		PrivateKey privateKey = decrypt
				.readPrivateKey(jksLocation, jksPassword);
		System.out.println("Ucitavanje setifikata...");
		// Certificate certificate = readRemoteCertificate(cerPath) ;
		Certificate certificate = readCerFromFile(cerPath);
		System.out.println("Reserializacija...");
		//doc = reserialize(doc);
		System.out.println("Potpisivanje...");
		signBody(body, privateKey, certificate);
		System.out.println("*****Potpisivanje*****");
	}

	public Certificate readCerFromFile(String path) {

		@SuppressWarnings("rawtypes")
		Collection collection;
		CertificateFactory certificateFactory;

		collection = null;
		certificateFactory = null;

		try {

			InputStream inputStream = this.getClass().getResourceAsStream(path);
			certificateFactory = CertificateFactory.getInstance("X.509");
			collection = certificateFactory.generateCertificates(inputStream);

			@SuppressWarnings("rawtypes")
			Iterator i = collection.iterator();
			while (i.hasNext()) {
				Certificate cert = (Certificate) i.next();
				return cert;
			}

		} catch (CertificateException e) {
			System.out.println(e.getMessage());
		} finally {

		}

		return null;
	}

	public void signBody(SOAPPart soap,PrivateKey privateKey,Certificate cert) {
		try {
			SOAPBody body = soap.getEnvelope().getBody();
			
			// kreira se signature objekat
			XMLSignature sig = new XMLSignature(soap, null,
					XMLSignature.ALGO_ID_SIGNATURE_RSA_SHA1);
			// kreiraju se transformacije nad dokumentom
			Transforms transforms = new Transforms(soap);
			
			XPathContainer cont = new XPathContainer(soap);
			
			cont.setXPath("//soap:Envelope/soap:Body/*");
			transforms.addTransform(Transforms.TRANSFORM_XPATH,
					cont.getElement());
			
			transforms.addTransform(Transforms.TRANSFORM_ENVELOPED_SIGNATURE);
			transforms.addTransform(Transforms.TRANSFORM_C14N_WITH_COMMENTS);
			
			sig.addDocument("", transforms, Constants.ALGO_ID_DIGEST_SHA1);
			
			// U KeyInfo se postavalja Javni kljuc samostalno i citav sertifikat
			sig.addKeyInfo(cert.getPublicKey());
			sig.addKeyInfo((X509Certificate)cert);
			
			body.appendChild(sig.getElement());
			
			// potpisivanje
			sig.sign(privateKey);
			
		} catch (Exception e) {
			throw new RuntimeException("signBody", e);
		}
		
	}
	
	
	public boolean verifyBody(SOAPPart soap) {
		try {
			SOAPBody body = soap.getEnvelope().getBody();
			
			// Pronalazi se prvi Signature element
			NodeList signatures = body.getElementsByTagNameNS(
					"http://www.w3.org/2000/09/xmldsig#", "Signature");
			
			if (signatures.getLength() == 0) {
				return false;
			}
			
			Element signatureEl = (Element) signatures.item(0);
			
			// kreira se signature objekat od elementa
			XMLSignature signature = new XMLSignature(signatureEl, null);
			// preuzima se key info
			KeyInfo keyInfo = signature.getKeyInfo();
			
			if (keyInfo != null) {
				keyInfo.registerInternalKeyResolver(new RSAKeyValueResolver());
				keyInfo.registerInternalKeyResolver(new X509CertificateResolver());
			}
			
			if (keyInfo == null || !keyInfo.containsX509Data()
					|| !keyInfo.itemX509Data(0).containsCertificate()) {
				return false;
			}
			
			X509Certificate cert = keyInfo.itemX509Data(0).itemCertificate(0)
					.getX509Certificate();
			
			if (!signature.checkSignatureValue(cert)) {
				return false;
			}
			
			
			body.removeChild(signatureEl);
			
			return true;
			
		} catch (Exception e) {
			System.out.println(e.getMessage());
			return false;
		}
	}
	
	
	public void print(Document doc) {
		
		
		try {
			Transformer tf = TransformerFactory.newInstance().newTransformer();
			
			tf.setOutputProperty(OutputKeys.INDENT, "yes");
			tf.setOutputProperty(
					"{http://xml.apache.org/xslt}indent-amount", "3");
			
			Source sc = new DOMSource(doc);
			
			StreamResult result = new StreamResult(System.out);
			System.out.println();
			tf.transform(sc, result);
			System.out.println();
			
		} catch (IllegalArgumentException | TransformerException e) {
			e.printStackTrace();
		}
		
	
	}

}
